import json
import math
import time
# 多线程
from threading import Thread

import cv2
import gym
import numpy as np
import win32con
import win32gui
import win32ui
from pykeyboard import PyKeyboard
from redis import Redis


# from pymouse import PyMouse


class HKEnv2D(gym.Env):
    def __init__(self):
        self.r = Redis(host='127.0.0.1', port=6379, db=1)
        self.hwnd_id = []
        self.k = PyKeyboard()
        self.initailize()
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)
        # 创建dc, 截取桌面
        self.hdesktop = win32gui.GetDesktopWindow()
        # 获取窗口大小
        l, t, r, b = win32gui.GetWindowRect(self.hwnd_id[0])
        self.width = r - l
        self.height = b - t

        # 创建设备描述表
        self.desktop_dc = win32gui.GetWindowDC(self.hdesktop)
        # 创建dc对象
        self.mfc_dc = win32ui.CreateDCFromHandle(self.desktop_dc)
        # 创建内存设备描述表
        self.save_dc = self.mfc_dc.CreateCompatibleDC()
        # 创建位图对象
        self.save_bit_map = win32ui.CreateBitmap()
        # 截图
        self.save_bit_map.CreateCompatibleBitmap(self.mfc_dc, self.width, self.height)

    def initailize(self):
        pass

    def get_all_hwnd(self, hwnd, mouse):
        if (win32gui.IsWindow(hwnd) and
            win32gui.IsWindowEnabled(hwnd) and
            win32gui.IsWindowVisible(hwnd)) and win32gui.GetWindowText(hwnd).__contains__('Hollow Knight'):
            self.hwnd_id.append(hwnd)

    def restart_game(self):
        win32gui.EnumWindows(self.get_all_hwnd, 0)
        time.sleep(1)
        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        time.sleep(1)
        self.k.release_key(self.k.left_key)
        self.k.release_key(self.k.right_key)
        self.k.release_key(self.k.up_key)
        self.k.release_key(self.k.down_key)
        self.k.release_key("Z")
        self.k.release_key("X")
        self.k.release_key("C")
        self.k.release_key("F")
        time.sleep(1)
        self.k.tap_key(self.k.left_key)
        time.sleep(4)
        self.k.tap_key(self.k.up_key)
        time.sleep(2)
        self.k.tap_key("z")

    def reset(self, *args, **kwargs):
        self._last_enemy_hp = 0
        self._last_knight_hp = 0
        self.s = time.time()
        self.e = time.time()
        t = Thread(target=self.restart_game)
        t.start()
        self.r.brpop(['start'])
        time.sleep(0.1)
        state_json = self.r.get('hollow_obs')
        binary = self.get_binary(save=True, no=0)
        self.state_to_obs(state_json)
        self._last_enemy_hp = self.enemy_hp
        self._last_knight_hp = self.knight_hp
        self.face = 'right'
        print("reset done")
        return binary

    def act(self, move, attack):
        self.do_attack(attack)
        self.do_move(move)

    def do_move(self, move):
        L = self.k.left_key
        R = self.k.right_key
        J = "Z"
        # 0: 无动作, 1: 跳跃, 2: 往左走, 3: 往右走, 4: 往左跳, 5 往右跳, 6 双跳
        if move == 0:
            # 释放左右和跳跃键
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key(J)
        if move == 1:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.press_key(J)
        if move == 2:
            self.k.release_key(R)
            self.k.release_key(J)
            self.k.press_key(L)
            self.face = 'left'
        if move == 3:
            self.k.release_key(L)
            self.k.release_key(J)
            self.k.press_key(R)
            self.face = 'right'
        if move == 4:
            self.k.release_key(R)
            self.k.press_key(L)
            self.k.press_key(J)
            self.face = 'left'
        if move == 5:
            self.k.release_key(L)
            self.k.press_key(R)
            self.k.press_key(J)
            self.face = 'right'
        if move == 6:
            self.k.release_key(J)
            self.k.press_key(J)

    def do_attack(self, attack):
        L = self.k.left_key
        R = self.k.right_key
        U = self.k.up_key
        D = self.k.down_key
        S = "F"
        C = "C"
        A = "X"
        # 0: 无动作, 1:左攻击, 2:右攻击, 3: 上攻击, 4: 下攻击, 5: 下砸, 6: 吼, 7: 左波, 8: 右波, 9: 左冲, 10: 右冲
        if attack in [1, 7, 9]:
            self.k.release_key(R)
            self.k.release_key(U)
            self.k.release_key(D)
            self.k.press_key(L)
            if self.face == 'right':
                time.sleep(0.04)
                self.face = 'left'
        if attack in [2, 8, 10]:
            self.k.release_key(L)
            self.k.release_key(U)
            self.k.release_key(D)
            self.k.press_key(R)
            if self.face == 'left':
                time.sleep(0.04)
            self.face = 'right'
        if attack in [3, 6]:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key(D)
            self.k.press_key(U)
            time.sleep(0.04)
        if attack in [4, 5]:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key(U)
            self.k.press_key(D)
            time.sleep(0.04)
        if attack in [1, 2, 3, 4]:
            self.k.tap_key(A)
        if attack in [5, 6, 7, 8]:
            self.k.tap_key(S)
        if attack in [9, 10]:
            self.k.tap_key(C)

    def get_binary(self, save=False, no=0):
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)

        self.save_dc.SelectObject(self.save_bit_map)
        # 截取从左上角(50, 50)长宽为(width, height)的图片
        # self.save_dc.BitBlt((0, 0), (self.width - 21, self.height - 41), self.mfc_dc, (0, 0), win32con.SRCCOPY)
        self.save_dc.BitBlt((0, 0), (self.width - 300, self.height - 210), self.mfc_dc, (150, 100), win32con.SRCCOPY)

        # 保存位图文件
        signedIntsArray = self.save_bit_map.GetBitmapBits(True)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.width, 4)
        # 删除多余部分
        img = img[0:self.height - 210, 0:self.width - 300]
        # 缩放
        img = cv2.resize(img, (200, 200))
        if save:
            cv2.imwrite("imgs/{}.png".format(no), img)
        # 保存
        # cv2.imwrite('test.png', img)
        # 转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if save:
            cv2.imwrite("imgs/gray_{}.png".format(no), gray)
        # 压缩到0-1
        gray = gray / 255.0
        # 缩放
        # gray = cv2.resize(gray, (511, 511))
        # 保存
        # cv2.imwrite('test_gray.png', gray)
        # 二值化
        # ret, binary = cv2.threshold(gray, 127, 1, cv2.THRESH_BINARY)
        # if save:
        #     ret, binary_255 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        #     cv2.imwrite("imgs/binary_{}.png".format(no), binary_255)
        # 转成numpy数组
        binary = np.array(gray)
        return binary

    def step(self, *args, **kwargs):
        # 获取状态
        reward_m = 0
        reward_a = 0
        # 如果boss在头顶
        if self.enemy_y > self.knight_y:
            # 下攻击
            if args[1] == 4:
                reward_a -= 1
        # 如果boss在脚下
        if self.enemy_y < self.knight_y:
            # 上攻击
            if args[1] == 3:
                reward_a -= 1
        # 如果boss在左边
        if self.enemy_x < self.knight_x:
            # 右攻击
            if args[1] == 2:
                reward_a -= 1
            # 右冲
            if args[1] == 10:
                reward_a -= 1
            # 右波
            if args[1] == 8:
                reward_a -= 1
        # 如果boss在右边
        if self.enemy_x > self.knight_x:
            # 左攻击
            if args[1] == 1:
                reward_a -= 1
            # 左冲
            if args[1] == 9:
                reward_a -= 1
            # 左波
            if args[1] == 7:
                reward_a -= 1
        if 0 < self.distance < 6:
            # 不攻击
            if args[1] == 0:
                reward_a -= 1
        if abs(self.enemy_x - self.knight_x) >= 3.2:
            if args[1] in [3, 4]:
                reward_a -= 1
            # 不移动
            if args[0] == 0:
                reward_m -= 1
            # boss在左边
            if self.enemy_x < self.knight_x:
                # 往左走
                if args[0] in [2]:
                    reward_m += 1
                # 往右走
                if args[0] in [3, 5]:
                    reward_m -= 1
                # 右冲
                if args[1] == 10:
                    reward_a -= 1
            # boss在右边
            if self.enemy_x > self.knight_x:
                # 往右走
                if args[0] in [3]:
                    reward_m += 1
                # 往左走
                if args[0] in [2, 4]:
                    reward_m -= 1
                # 左冲
                if args[1] == 9:
                    reward_a -= 1
        if abs(self.enemy_x - self.knight_x) >= 6:
            # boss在右边
            if self.enemy_x > self.knight_x:
                # 右冲
                if args[1] == 10:
                    reward_a += 0.5
            # boss在左边
            if self.enemy_x < self.knight_x:
                # 左冲
                if args[1] == 9:
                    reward_a += 0.5
        if self.distance <= 1.5:
            if args[0] == 0:
                reward_m -= 1
            # boss在右边
            if self.enemy_x > self.knight_x:
                # 右冲
                if args[1] == 10:
                    reward_a += 1
                # 右移
                if args[0] in [5, 3]:
                    reward_m -= 1
            # boss在左边
            if self.enemy_x < self.knight_x:
                # 左冲
                if args[1] == 9:
                    reward_a += 1
                # 左移
                if args[0] in [2, 4]:
                    reward_m -= 1
        start = time.time()
        self.act(args[0], args[1])
        # self.act(0, 2)
        max_time = 0.07
        # if args[1] in [9, 10]:
        #     max_time = 0.05
        while True:
            end = time.time()
            state_json = self.r.get('hollow_obs')
            self.state_to_obs(state_json)
            if self.is_attacking or self.knight_spell or self.dashing or end - start < max_time:
                time.sleep(0.001)
                continue
            if end - start < max_time + 0.02:
                continue
            else:
                break
        duration = end - start
        if args[1] in [9, 10]:
            self.is_sprint = False
        elif args[1] != 0:
            self.is_sprint = True
        binary = self.get_binary()
        state_json = self.r.get('hollow_obs')
        self.state_to_obs(state_json)
        # if 2 <= abs(self.enemy_x - self.knight_x) <= 5.5 and self.distance < 5:
        #     if args[1] in [9, 10]:
        #         reward_a += 0.5
        if self.distance < 2:
            if not args[1] in [9, 10]:
                reward_m -= 1
            else:
                reward_a -= 1
        if abs(self.enemy_x - self.knight_x) > 6:
            if args[1] in [9, 10]:
                reward_a -= 1
        if 3 <= self.distance <= 5:
            if not args in [9, 10]:
                reward_m += 1
        self.done = False
        self.key_point = False
        if self._last_enemy_hp > self.enemy_hp:
            self.key_point = True
            reward_a = (self._last_enemy_hp - self.enemy_hp) / 20 / duration * 0.2
            reward_m = reward_a
            # 如果下攻击
            if args[1] == 4:
                reward_a -= 0.5
            if args[1] in [5, 6, 7, 8]:
                reward_a += 0.5
            reward_a += 10
            reward_m += 10
        if self._last_knight_hp > self.knight_hp:
            self.key_point = True
            if args[1] in [5, 6, 7, 8, 9, 10]:
                reward_a = -10
            else:
                reward_a = -10
                reward_m = -10
        if self.knight_hp <= 0:
            self.done = True
        if str(self.r.get('done')) == "1":
            self.done = True
            reward_a = 5
            reward_m = 5
        self._last_enemy_hp = self.enemy_hp
        self._last_knight_hp = self.knight_hp
        if self.done:
            self.get_binary(save=True, no=1)
            self.k.release_key(self.k.left_key)
            self.k.release_key(self.k.right_key)
            self.k.release_key(self.k.up_key)
            self.k.release_key(self.k.down_key)
            self.k.release_key("Z")
            self.k.release_key("X")
            self.k.release_key("C")
            self.k.release_key("F")
        return binary, reward_m, reward_a, self.done, self.key_point

    def state_to_obs(self, state):
        state_dict = json.loads(state)
        # 是否黑冲
        self.shadow_dashing = True if state_dict['shadow_dashing'] == '1' else False

        # 是否冲刺
        self.dashing = True if state_dict['dashing'] == '1' else False

        # 是否在攻击
        self.is_attacking = True if state_dict['is_attacking'] == '1' else False

        # 是否在释放技能
        self.knight_spell = True if state_dict['knight_spell'] != "Inactive" else False

        # 是否在地面
        self.on_ground = True if state_dict['onGround'] == '1' else False
        # 血量
        self.knight_hp = int(state_dict['knight_hp'])
        self.enemy_hp = int(state_dict['enemy_hp'])

        # knight 的位置
        self.knight_x = float(state_dict['knight_x'])
        self.knight_y = float(state_dict['knight_y'])

        # knight的灵魂量
        self.knight_mp = float(state_dict['knight_mp'])

        # boss 的位置
        self.enemy_x = float(state_dict['enemy_x'])
        self.enemy_y = float(state_dict['enemy_y'])

        # 距离, 用于判断是否在攻击范围内, 平方和开方,
        self.distance = math.sqrt((self.knight_x - self.enemy_x) ** 2 + (self.knight_y - self.enemy_y) ** 2)

        # 冲刺是否冷却
        self.is_sprint = True


if __name__ == '__main__':
    env = HKEnv2D()
    # env.restart_game()z
    win32gui.EnumWindows(env.get_all_hwnd, 0)
    handle = env.hwnd_id[0]
    # 激活窗口
    win32gui.SetForegroundWindow(handle)
    time.sleep(3)
    k = PyKeyboard()
    k.press_key("Z")
    time.sleep(0.2)
    k.release_key("Z")
    time.sleep(0.1)
    k.press_key("Z")
    time.sleep(1)
    k.release_key('Z')

    obs = env.reset()
    print(obs)
    for i in range(100):
        action = np.random.randint(0, 12)
        obs, reward, done, info = env.step(1, 1)
        count = 0
        print(obs, reward, done, info)
